package cn.xiaopengstack.config;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.springframework.ai.chat.model.ChatModel;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author jiangyangang
 */
@Data
@AllArgsConstructor
@NoArgsConstructor
public class ChatModelRegistry{
    private ConcurrentHashMap<ModelNameVal, ChatModel> registry;

    public ChatModel getModel(String displayModelName){
        for (Map.Entry<ModelNameVal, ChatModel> entry : registry.entrySet()) {
            ModelNameVal modelNameVal = entry.getKey();
            if (!modelNameVal.getDisplayName().equals(displayModelName)) {
                continue;
            }
            return entry.getValue();
        }

        throw new RuntimeException("未知的模型[" + displayModelName +"]");
    }

    public String getOriginModelName(String displayModelName){
        for (Map.Entry<ModelNameVal, ChatModel> entry : registry.entrySet()) {
            ModelNameVal modelNameVal = entry.getKey();
            if (!modelNameVal.getDisplayName().equals(displayModelName)) {
                continue;
            }
            return entry.getKey().getModelRequireName();
        }
        return null;
    }

    public void register(ModelNameVal modelNameVal, ChatModel chatModel){
        registry.put(modelNameVal, chatModel);
    }
}